#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from rclpy.time import Time
from rclpy.action import ActionClient
from rclpy.executors import MultiThreadedExecutor
from rclpy.callback_groups import ReentrantCallbackGroup, MutuallyExclusiveCallbackGroup

from geometry_msgs.msg import Twist, Point
from std_msgs.msg import String, Bool
from nav_msgs.msg import Odometry
from ball_collector_interfaces.msg import BallInfo, RobotState
from ball_collector_interfaces.srv import SetGripper, ChangeColor
import threading
import time
import math
import random
import argparse
from enum import Enum


class TestMode(Enum):
    """测试模式枚举"""
    ALL = 'all'               # 完整测试所有状态
    SEARCH = 'search'         # 只测试搜索状态
    APPROACH = 'approach'     # 只测试接近状态
    ALIGN = 'align'           # 只测试对准状态
    GRASP = 'grasp'           # 只测试抓取状态
    TRANSPORT = 'transport'   # 只测试运输状态
    RELEASE = 'release'       # 只测试释放状态
    RETURN = 'return'         # 只测试返回状态
    LOST = 'lost'             # 测试球丢失处理
    COLOR = 'color'           # 测试颜色切换


class TestStateMonitor(Node):
    """
    监控策略节点状态的测试节点
    """
    def __init__(self, test_mode='all'):
        super().__init__('test_state_monitor')
        
        # 设置测试模式
        self.test_mode = TestMode(test_mode)
        self.get_logger().info(f"测试模式: {self.test_mode.value}")
        
        # 创建互斥回调组
        self.timer_cb_group = MutuallyExclusiveCallbackGroup()
        self.topic_cb_group = MutuallyExclusiveCallbackGroup()
        
        # 状态追踪
        self.current_state = "IDLE"
        self.has_ball = False
        self.cmd_vel_history = []
        self.state_history = []
        self.gripper_calls = []
        self.target_color = "red"
        
        # 标记测试进度
        self.test_step = 0
        self.step_start_time = time.time()
        self.state_change_time = time.time()
        self.test_completed = False
        
        # 订阅策略节点的输出
        self.cmd_vel_sub = self.create_subscription(
            Twist,
            'cmd_vel',
            self.cmd_vel_callback,
            10,
            callback_group=self.topic_cb_group
        )
        
        self.robot_state_sub = self.create_subscription(
            RobotState,
            'robot_state',
            self.robot_state_callback,
            10,
            callback_group=self.topic_cb_group
        )
        
        # 创建发布者
        self.ball_pub = self.create_publisher(
            BallInfo,
            'detected_ball',
            10
        )
        
        self.odom_pub = self.create_publisher(
            Odometry,
            'odom',
            10
        )
        
        # 创建模拟服务
        self.gripper_service = self.create_service(
            SetGripper,
            '/hardware_interface/set_gripper',
            self.gripper_callback
        )
        
        self.color_service = self.create_service(
            ChangeColor,
            'change_target_color',
            self.color_callback
        )
        
        # 创建定时器
        self.test_timer = self.create_timer(
            0.1,  # 100ms, 更高频率以确保响应更迅速
            self.test_step_callback,
            callback_group=self.timer_cb_group
        )
        
        self.status_timer = self.create_timer(
            5.0,  # 5秒
            self.print_status,
            callback_group=self.timer_cb_group
        )
        
        # 初始化测试参数
        self.ball_visible = False
        self.ball_x = 320
        self.ball_y = 240
        self.ball_diameter = 100
        
        self.get_logger().info("测试监控节点已初始化")
    
    def cmd_vel_callback(self, msg):
        """记录速度命令"""
        current_time = time.time()
        cmd = {
            'time': current_time,
            'linear_x': msg.linear.x,
            'angular_z': msg.angular.z
        }
        self.cmd_vel_history.append(cmd)
        
        # 仅记录有意义的变化
        if len(self.cmd_vel_history) > 1:
            prev = self.cmd_vel_history[-2]
            if (abs(prev['linear_x'] - cmd['linear_x']) > 0.01 or 
                abs(prev['angular_z'] - cmd['angular_z']) > 0.01):
                self.get_logger().info(f"速度命令: 线速度={msg.linear.x:.2f}, 角速度={msg.angular.z:.2f}")
    
    def robot_state_callback(self, msg):
        """记录机器人状态"""
        if msg.current_state != self.current_state or msg.has_ball != self.has_ball:
            self.get_logger().info(f"状态变化: {self.current_state} -> {msg.current_state}, 持球: {msg.has_ball}")
            
            # 记录状态变化时间
            self.state_change_time = time.time()
            
            # 记录状态历史
            self.state_history.append({
                'time': time.time(),
                'state': msg.current_state,
                'has_ball': msg.has_ball
            })
            
            self.current_state = msg.current_state
            self.has_ball = msg.has_ball
            self.target_color = msg.target_color
            
            # 根据状态变化重置测试步骤
            if self.test_mode == TestMode.ALL:
                if msg.current_state == "IDLE" and self.test_step == 0:
                    self.test_step = 1
                elif msg.current_state == "SEARCHING" and self.test_step == 1:
                    self.test_step = 2
                elif msg.current_state == "APPROACHING" and self.test_step == 2:
                    self.test_step = 3
                elif msg.current_state == "ALIGNING" and self.test_step == 3:
                    self.test_step = 4
                elif msg.current_state == "GRASPING" and self.test_step == 4:
                    self.test_step = 5
                elif msg.current_state == "TRANSPORTING" and self.test_step == 5:
                    self.test_step = 6
                elif msg.current_state == "RELEASING" and self.test_step == 6:
                    self.test_step = 7
                elif msg.current_state == "RETURNING" and self.test_step == 7:
                    self.test_step = 8
                elif msg.current_state == "SEARCHING" and self.test_step == 8:
                    # 回到搜索状态，测试完成
                    self.get_logger().info("完成了一个完整的收集周期!")
                    
    def gripper_callback(self, request, response):
        """模拟夹爪服务回调"""
        self.get_logger().info(f"夹爪调用: grip={request.grip}, force={request.force}")
        
        # 记录夹爪调用
        self.gripper_calls.append({
            'time': time.time(),
            'grip': request.grip,
            'force': request.force
        })
        
        # 始终返回成功
        response.success = True
        response.message = f"夹爪{'关闭' if request.grip else '打开'}成功"
        return response
    
    def color_callback(self, request, response):
        """模拟颜色更改服务回调"""
        self.get_logger().info(f"颜色更改请求: {request.new_color}")
        
        # 记录目标颜色变化
        self.target_color = request.new_color
        
        # 返回成功
        response.success = True
        response.current_color = request.new_color
        return response
    
    def publish_mock_ball(self, is_visible=True, x=320, y=240, diameter=100, color=None):
        """发布模拟球体检测信息"""
        if not is_visible:
            # 不发布，相当于没有检测到球
            return
        
        msg = BallInfo()
        msg.header.stamp = self.get_clock().now().to_msg()
        msg.header.frame_id = "camera_link"
        msg.id = 1
        msg.color = color if color else self.target_color
        
        msg.center = Point()
        msg.center.x = float(x)
        msg.center.y = float(y)
        msg.center.z = 0.0
        
        msg.diameter = float(diameter)
        msg.confidence = 0.95
        msg.is_target = True
        
        self.ball_pub.publish(msg)
        self.get_logger().debug(f"发布球信息: x={x}, y={y}, 直径={diameter}, 颜色={msg.color}")
    
    def publish_mock_odometry(self, x=0.0, y=0.0, theta=0.0):
        """发布模拟里程计数据"""
        msg = Odometry()
        msg.header.stamp = self.get_clock().now().to_msg()
        msg.header.frame_id = "odom"
        msg.child_frame_id = "base_link"
        
        # 位置
        msg.pose.pose.position.x = x
        msg.pose.pose.position.y = y
        msg.pose.pose.position.z = 0.0
        
        # 将偏航角转换为四元数（简化版本）
        msg.pose.pose.orientation.w = math.cos(theta / 2.0)
        msg.pose.pose.orientation.z = math.sin(theta / 2.0)
        
        self.odom_pub.publish(msg)
    
    def test_step_callback(self):
        """测试步骤回调，根据测试模式执行不同测试"""
        if self.test_completed:
            return
        
        # 始终发布里程计数据
        self.publish_mock_odometry(0.0, 0.0, 0.0)
        
        # 根据测试模式执行不同测试
        if self.test_mode == TestMode.ALL:
            self.run_all_tests()
        elif self.test_mode == TestMode.SEARCH:
            self.test_searching()
        elif self.test_mode == TestMode.APPROACH:
            self.test_approaching()
        elif self.test_mode == TestMode.ALIGN:
            self.test_aligning()
        elif self.test_mode == TestMode.GRASP:
            self.test_grasping()
        elif self.test_mode == TestMode.TRANSPORT:
            self.test_transporting()
        elif self.test_mode == TestMode.RELEASE:
            self.test_releasing()
        elif self.test_mode == TestMode.RETURN:
            self.test_returning()
        elif self.test_mode == TestMode.LOST:
            self.test_ball_lost()
        elif self.test_mode == TestMode.COLOR:
            self.test_color_change()

    def run_all_tests(self):
        """运行完整测试序列"""
        current_time = time.time()
        elapsed_since_state_change = current_time - self.state_change_time
        
        if self.test_step == 0:
            # 初始化阶段 - 等待策略节点启动
            self.get_logger().info("等待策略节点启动...")
            if self.current_state == "IDLE":
                self.test_step = 1
                self.step_start_time = current_time
        
        elif self.test_step == 2 and self.current_state == "SEARCHING":
            # 已进入搜索状态，等待一段时间后发送球信息
            if elapsed_since_state_change > 3.0:
                self.get_logger().info("发送球信息，触发接近行为")
                self.publish_mock_ball(True, 320, 240, 80)
        
        elif self.test_step == 3 and self.current_state == "APPROACHING":
            # 接近状态 - 逐渐增大球直径模拟靠近
            elapsed = current_time - self.step_start_time
            ball_diameter = min(192, 80 + elapsed * 10)  # 逐渐增大但不超过192
            self.publish_mock_ball(True, 320, 240, ball_diameter)
        
        elif self.test_step == 4 and self.current_state == "ALIGNING":
            # 对准状态 - 从偏离中心逐渐移动到中心
            elapsed = current_time - self.step_start_time
            x_pos = max(320, 420 - elapsed * 20)  # 从右侧逐渐移动到中心
            self.publish_mock_ball(True, x_pos, 240, 192)  # 大直径表示已经很近
        
        elif self.test_step == 5 and self.current_state == "GRASPING":
            # 抓取状态 - 保持球在中心位置
            self.publish_mock_ball(True, 320, 240, 192)
        
        elif self.test_step == 6 and self.current_state == "TRANSPORTING":
            # 运输状态 - 不发送球信息，模拟已抓取
            # 模拟移动到目标区域
            elapsed = current_time - self.step_start_time
            x = min(2.0, 1.0 + elapsed * 0.2)
            y = min(1.0, elapsed * 0.2)
            self.publish_mock_odometry(x, y, 0.0)
        
        elif self.test_step == 7 and self.current_state == "RELEASING":
            # 释放状态 - 保持在目标位置
            self.publish_mock_odometry(2.0, 1.0, 0.0)
        
        elif self.test_step == 8 and self.current_state == "RETURNING":
            # 返回状态 - 模拟返回
            elapsed = current_time - self.step_start_time
            x = max(0.5, 2.0 - elapsed * 0.3)
            y = max(0.0, 1.0 - elapsed * 0.2)
            self.publish_mock_odometry(x, y, elapsed * 0.1)
        
        elif self.test_step >= 9:
            # 测试完成
            if not self.test_completed:
                self.get_logger().info("全部测试完成!")
                self.test_completed = True
                self.analyze_results()

    def test_searching(self):
        """测试搜索状态"""
        # 发布里程计数据，模拟机器人原地旋转
        current_time = time.time()
        elapsed = current_time - self.step_start_time
        self.publish_mock_odometry(0.0, 0.0, math.sin(elapsed * 0.5))
        
        # 每隔5秒发送一次球信息，测试搜索-接近循环
        if int(elapsed) % 10 < 5:
            # 不发送球信息，保持搜索
            pass
        else:
            # 发送球信息，触发接近行为
            self.publish_mock_ball(True, 320, 240, 80)
    
    def test_approaching(self):
        """测试接近状态"""
        # 确保处于SEARCHING或APPROACHING状态
        if self.current_state not in ["SEARCHING", "APPROACHING"]:
            self.get_logger().info("等待进入搜索或接近状态...")
            return
            
        # 发布固定位置但直径逐渐增大的球，模拟接近
        current_time = time.time()
        elapsed = current_time - self.step_start_time
        ball_diameter = min(200, 50 + (elapsed % 10) * 15)  # 直径在50-200之间循环变化
        
        self.publish_mock_ball(True, 320, 240, ball_diameter)
        self.publish_mock_odometry(0.0, 0.0, 0.0)
    
    def test_aligning(self):
        """测试对准状态"""
        # 确保处于接近或对准状态
        if self.current_state not in ["APPROACHING", "ALIGNING"]:
            # 如果在搜索状态，发送足够大的球触发接近
            if self.current_state == "SEARCHING":
                self.publish_mock_ball(True, 340, 240, 150)
            return
            
        # 发布位置不断变化的球，模拟对准过程
        current_time = time.time()
        elapsed = current_time - self.step_start_time
        
        # 球在中心附近左右移动
        x_center = 320 + 100 * math.sin(elapsed)
        
        self.publish_mock_ball(True, x_center, 240, 180)
        self.publish_mock_odometry(0.5, 0.0, 0.0)
    
    def test_grasping(self):
        """测试抓取状态"""
        # 确保进入对准或抓取状态
        if self.current_state not in ["ALIGNING", "GRASPING"]:
            # 如果在其他状态，尝试让机器人进入对准状态
            if self.current_state == "SEARCHING":
                self.publish_mock_ball(True, 320, 240, 150)
            elif self.current_state == "APPROACHING":
                self.publish_mock_ball(True, 320, 240, 192)
            return
            
        # 发布完全对准的球
        self.publish_mock_ball(True, 320, 240, 200)
        self.publish_mock_odometry(1.0, 0.0, 0.0)
    
    def test_transporting(self):
        """测试运输状态"""
        # 确保处于运输状态或已有球
        if not self.has_ball and self.current_state != "TRANSPORTING":
            # 如果没有球并且不在运输状态，先测试抓取
            self.test_grasping()
            return
            
        # 模拟机器人移动到目标区域
        current_time = time.time()
        elapsed = current_time - self.step_start_time
        
        # 计算机器人位置，模拟向目标区域移动
        x = min(2.0, (elapsed % 20) * 0.1)
        y = min(1.0, (elapsed % 10) * 0.1)
        
        self.publish_mock_odometry(x, y, 0.0)
    
    def test_releasing(self):
        """测试释放状态"""
        # 确保处于释放状态或运输状态
        if self.current_state not in ["TRANSPORTING", "RELEASING"]:
            # 如果不在正确状态，尝试让机器人进入运输状态
            self.test_transporting()
            return
            
        # 模拟已到达目标区域
        self.publish_mock_odometry(2.0, 1.0, 0.0)
    
    def test_returning(self):
        """测试返回状态"""
        # 确保处于返回状态
        if self.current_state != "RETURNING":
            # 如果不在返回状态，无法直接测试，提示信息
            self.get_logger().info("无法直接进入返回状态，请先完成球的运输和释放")
            # 可以进行完整测试循环
            self.test_mode = TestMode.ALL
            return
            
        # 模拟从目标区域返回
        current_time = time.time()
        elapsed = current_time - self.step_start_time
        
        x = max(0.0, 2.0 - (elapsed % 10) * 0.2)
        y = max(0.0, 1.0 - (elapsed % 5) * 0.2)
        
        self.publish_mock_odometry(x, y, elapsed * 0.1)
    
    def test_ball_lost(self):
        """测试球丢失行为"""
        current_time = time.time()
        elapsed = current_time - self.step_start_time
        
        # 周期性地显示和隐藏球
        if int(elapsed * 2) % 4 == 0:  # 25%的时间可见
            self.publish_mock_ball(True, 320, 240, 100)
            self.get_logger().debug("发布球信息 - 可见")
        else:
            # 不发布球信息，模拟丢失
            self.get_logger().debug("不发布球信息 - 丢失")
        
        self.publish_mock_odometry(1.0, 0.5, 0.0)
    
    def test_color_change(self):
        """测试颜色切换"""
        current_time = time.time()
        elapsed = current_time - self.step_start_time
        
        # 循环测试不同颜色
        colors = ["red", "green", "blue", "yellow"]
        color_index = int(elapsed) % len(colors)
        current_color = colors[color_index]
        
        # 每3秒切换一次颜色
        if int(elapsed) != int(elapsed - 0.1) and int(elapsed) % 3 == 0:
            self.get_logger().info(f"测试颜色切换到: {current_color}")
            # 如果颜色已经变化，可以直接发布对应颜色的球
            self.publish_mock_ball(True, 320, 240, 100, current_color)
        else:
            # 使用当前目标颜色发布球
            self.publish_mock_ball(True, 320, 240, 100)
        
        self.publish_mock_odometry(0.0, 0.0, 0.0)
    
    def print_status(self):
        """打印当前测试状态"""
        if not self.test_completed:
            self.get_logger().info(f"测试模式: {self.test_mode.value}, 当前状态: {self.current_state}, 持球: {self.has_ball}")
            
            # 打印最后的命令
            if self.cmd_vel_history:
                last_cmd = self.cmd_vel_history[-1]
                self.get_logger().info(f"最后的速度命令: 线速度={last_cmd['linear_x']:.2f}, 角速度={last_cmd['angular_z']:.2f}")
    
    def analyze_results(self):
        """分析测试结果"""
        self.get_logger().info("测试结果分析:")
        
        # 分析状态转换
        if len(self.state_history) > 1:
            self.get_logger().info("状态转换序列:")
            for i, state in enumerate(self.state_history):
                transition_time = state['time']
                if i > 0:
                    transition_time = state['time'] - self.state_history[i-1]['time']
                self.get_logger().info(f"  {i}: {state['state']} (持球: {state['has_ball']}), 花费时间: {transition_time:.2f}s")
        
        # 分析夹爪调用
        if self.gripper_calls:
            self.get_logger().info("夹爪操作序列:")
            for i, call in enumerate(self.gripper_calls):
                self.get_logger().info(f"  {i}: {'关闭' if call['grip'] else '打开'}, 力度: {call['force']}")
        
        # 检查是否完成了完整的收集周期
        states_seen = set(s['state'] for s in self.state_history)
        expected_states = {"SEARCHING", "APPROACHING", "ALIGNING", "GRASPING", 
                         "TRANSPORTING", "RELEASING", "RETURNING"}
        
        missing_states = expected_states - states_seen
        if missing_states:
            self.get_logger().warn(f"测试未覆盖以下状态: {missing_states}")
        else:
            self.get_logger().info("测试覆盖了所有期望的状态转换")
        
        # 分析是否检测到了球的抓取和释放
        ball_pickup = False
        ball_release = False
        for i in range(1, len(self.state_history)):
            prev = self.state_history[i-1]
            curr = self.state_history[i]
            if not prev['has_ball'] and curr['has_ball']:
                ball_pickup = True
            if prev['has_ball'] and not curr['has_ball']:
                ball_release = True
        
        if ball_pickup:
            self.get_logger().info("✓ 成功检测到球的抓取")
        else:
            self.get_logger().warn("✗ 未检测到球的抓取")
            
        if ball_release:
            self.get_logger().info("✓ 成功检测到球的释放")
        else:
            self.get_logger().warn("✗ 未检测到球的释放")


def main(args=None):
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='策略节点测试')
    parser.add_argument('--test', type=str, default='all', 
                    choices=['all', 'search', 'approach', 'align', 'grasp', 
                             'transport', 'release', 'return', 'lost', 'color'],
                    help='指定要运行的测试模式')
    parser.add_argument('--debug', action='store_true',
                    help='启用调试日志')
    
    parsed_args, remaining_args = parser.parse_known_args(args)
    
    rclpy.init(args=remaining_args)
    
    # 创建测试节点
    test_node = TestStateMonitor(test_mode=parsed_args.test)
    
    # 设置日志级别
    if parsed_args.debug:
        rclpy.get_global_logging().set_logger_level(test_node.get_name(), rclpy.logging.LoggingSeverity.DEBUG)
    
    # 创建多线程执行器
    executor = MultiThreadedExecutor(num_threads=4)
    executor.add_node(test_node)
    
    # 启动执行器
    executor_thread = threading.Thread(target=executor.spin, daemon=True)
    executor_thread.start()
    
    try:
        # 等待测试完成或超时
        timeout = 180  # 3分钟超时
        start_time = time.time()
        
        while not test_node.test_completed and time.time() - start_time < timeout:
            time.sleep(0.1)
        
        if not test_node.test_completed and time.time() - start_time >= timeout:
            test_node.get_logger().warn(f"测试超时 ({timeout}秒)")
            test_node.analyze_results()
    
    except KeyboardInterrupt:
        test_node.get_logger().info("测试被用户中断")
        test_node.analyze_results()
    
    finally:
        executor.shutdown()
        test_node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()